# from scipy.stats import pearsonr
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from io import BytesIO
import base64

# 该模块用于特征和标签间的皮尔逊相关系数分析
class Pearson():
    df_features = None
    pearsonMatrix = []
    df_corr_matrix = None
    def calculatePearson(self,df_features):
        '''
        计算特征间的皮尔逊相关系数
        :param df_featuresAndTarget: (pd.DataFrame()) 包含特征数据的DataFrame表格
        :return: 皮尔逊系数矩阵 DataFrame
        '''
        self.df_features = df_features
        df_corr_matrix = df_features.corr()
        self.df_corr_matrix = df_corr_matrix
        return df_corr_matrix
    def classificationByThreshold(self,Threshold,df_corr_matrix = None):
        '''
        通过阈值将相关性大的特征分在一起
        :param Threshold: 皮尔逊系数阈值
        :param df_corr_matrix: 皮尔逊矩阵DataFrame 默认使用calculatePearson()方法的返回值
        :return: [[相关性大的特征名称],[]]
        '''
        if self.df_corr_matrix is not None and df_corr_matrix is None: df_corr_matrix = self.df_corr_matrix
        upper = df_corr_matrix.where(np.triu(np.ones(df_corr_matrix .shape), k=1).astype(np.bool))
        to_drop = [column for column in upper.columns if any(abs(upper[column]) > Threshold)]
        rest = self.df_features.drop(to_drop, axis=1)

        # 创建一个空列表存放相关系数大的特征对
        corr_features_list = []
        # 创建一个临时列表存放每次
        for i in range(len(rest.columns)):  # 遍历剩余的特征名，找和它相关性强的特征存在一个列表中
            # 当前行名
            row = rest.columns[i]
            corr_features = [row]
            # 当前特征名在upper中的行索引值
            row_index = list(upper.index).index(rest.columns[i])
            #     print(row_index)
            for j in range(i + 1, len(upper.columns)):  # 遍历相关矩阵的列，找相关性强的特征
                corr_value = upper.iloc[row_index, j]
                if abs(corr_value) > Threshold:
                    corr_features.append(upper.columns[j])
            corr_features_list.append(corr_features)
        return corr_features_list

    def plotHeatMap(self,df_features):
        corr_mat = df_features.corr()
        f, ax = plt.subplots(figsize=(7, 7))
        mask = np.zeros_like(corr_mat)
        for i in range(1, len(mask)):
            for j in range(0, i):
                mask[j][i] = True
        sns.heatmap(corr_mat, annot=True, mask=mask, linewidths=.05, square=True, annot_kws={'size': 12})
        # print(corr_mat)
        plt.xticks(size=8, rotation=0)
        plt.yticks(size=8, rotation=0)
        plt.subplots_adjust(left=.1, right=0.95, bottom=0.22, top=0.95)
        # colorbar 刻度线设置
        cax = plt.gcf().axes[-1]
        cax.tick_params(labelsize=18)  # colorbar 刻度字体大小
        # plt.tight_layout()
        # plt.gcf().subplots_adjust(left=0.05,top=0.91,bottom=0.09)
        # plt.show()

        # 转base64 方便传输给前端
        figfile = BytesIO()
        plt.savefig(figfile, format='png')
        figfile.seek(0)
        figdata_png = base64.b64encode(figfile.getvalue())  # 将图片转为base64
        figdata_str = str(figdata_png, "utf-8")  # 提取base64的字符串，不然是b'xxx'

        # 保存为.html
        html = '<img src=\"data:image/png;base64,{}\"/>'.format(figdata_str)
        return  html





# df1 = pd.read_excel(r'D:\Jupyter code\物理特征提升带隙预测精度\第一次修改\表数据\87物理特征训练集.xlsx')
# targetName = 'Perovskite_band_gap'
# featureName = ['IR_t', 'IR_u', 'IR_A', 'IR_B', 'IR_X', 'IR_sumAX', 'IR_sumBX',
#        'IR_diffAX', 'IR_diffBX', 'IR_diffAB', 'IRquoAX', 'IRquoBX', 'EN_B',
#        'EN_X', 'EN_diffBX', 'electron_affinity_B', 'electron_affinity_X',
#        'electron_affinity_diffBX', 'dipole_polarizability_B',
#        'dipole_polarizability_X', 'dipole_polarizability_diffBX',
#        'atomic_radius_B', 'atomic_radius_X', 'atomic_radius_diffBX', 'd_X',
#        'f_B',]
# features=df1.loc[:,featureName]
# target=df1['Perovskite_band_gap']
#
# Pearson=Pearson()
# Pearson.calculatePearson(features)
# corr_features_list = Pearson.classificationByThreshold(0.95)
# print(corr_features_list)